class Env:
    def __init__(self):
        self._states = set()
        self._state = None
        self._actions = []
        self._gamma = None

    @property
    def states(self):
        return self._states

    @property
    def state_space(self):
        return self._state_shape

    @property
    def actions(self):
        return self._actions

    @property
    def action_space(self):
        return len(self._actions)

    @property
    def gamma(self):
        return self._gamma

    def _world_init(self):
        raise NotImplementedError

    def reset(self):
        raise NotImplementedError

    def step(self, state, action):
        """Return distribution and next states"""
        raise NotImplementedError

    def set_state(self, state):
        self._state = state


class MatrixEnv(Env):
    def __init__(self, height=4, width=4):
        super().__init__()

        self._action_space = 4
        self._actions = list(range(4))

        self._state_shape = (2,)
        self._state_shape = (height, width)
        self._states = [(i, j) for i in range(height) for j in range(width)]

        self._gamma = 0.9
        self._height = height
        self._width = width

        self._world_init()

    @property
    def state(self):
        return self._state

    @property
    def gamma(self):
        return self._gamma

    def set_gamma(self, value):
        self._gamma = value

    def reset(self):
        self._state = self._start_point

    def _world_init(self):
        # start_point
        self._start_point = (0, 0)
        self._end_point = (self._height - 1, self._width - 1)

    def _state_switch(self, act):
        # 0: h - 1, 1: w + 1, 2: h + 1, 3: w - 1
        if act == 0:  # up
            self._state = (max(0, self._state[0] - 1), self._state[1])
        elif act == 1:  # right
            self._state = (self._state[0], min(self._width - 1, self._state[1] + 1))
        elif act == 2:  # down
            self._state = (min(self._height - 1, self._state[0] + 1), self._state[1])
        elif act == 3:  # left
            self._state = (self._state[0], max(0, self._state[1] - 1))

    def step(self, act):
        assert 0 <= act <= 3

        done = False
        reward = 0.

        self._state_switch(act)

        if self._state == self._end_point:
            reward = 1.
            done = True

        return None, done, [reward], [self._state]